#!/usr/bin/env python3
# Author: Armit
# Create Time: 2022/11/18 

from collections import Counter

import numpy as np
import matplotlib.pyplot as plt

from data import *


def get_dist(p1: Vector3d, p2: Vector3d) -> float:
  return ((np.asarray(p1) - np.asarray(p2)) ** 2).sum() ** 0.5

def preprocess():
  residue_names, sequences, locations = load_pdb_data()
  
  n_cfm = len(sequences)
  print(f'>> n_cfm: {n_cfm}')

  for seq, loc in zip(sequences, locations):
    n_seq = len(seq)
    print('>> len_seq:', n_seq)
    loc = np.asarray(loc)   # [N, D=3]

    if not 'plot 3d':
      plt.clf()
      ax = plt.axes(projection='3d')
      ax.scatter(loc[:, 0], loc[:, 1], zs=loc[:, 2], zdir='z', s=40)
      plt.tight_layout()
      plt.show()

    # 计算平均链上近邻距离
    dists = [get_dist(loc[i], loc[i-1]) for i in range(1, n_seq)]
    avg_rr_dist = sum(dists) / len(dists)
    # 设定探测距离阈值
    prob_r = avg_rr_dist * 1.2

    # 收集非共价的近邻关系
    pair_dist = []
    def test_and_collect(i, j):
      d = get_dist(loc[i], loc[j])
      if d > prob_r: return

      r_r = '-'.join(sorted([seq[i], seq[j]]))
      pair_dist.append((r_r, d))

    for i in range(1, len(seq)-1):
      for j in range(0, i-2):
        test_and_collect(i, j)
      for j in range(i+2, n_seq):
        test_and_collect(i, j)

    n_pairs = len(pair_dist)
    print(f'>> n_pairs: {n_pairs}')

    if 'plot hist dists':
      pairs = [p for p, d in pair_dist]
      cnt_pair = sorted([(c, p) for p, c in Counter(pairs).items()], reverse=True)
      count = [c for c, p in cnt_pair]
      pair  = [p for c, p in cnt_pair]
      print(pair)

      dists = [d for p, d in pair_dist]
      plt.subplot(211)
      plt.hist(dists, bins=32)
      plt.subplot(212)
      plt.plot(count)
      plt.show()


if __name__ == '__main__':
  preprocess()
